# 图像分类
import torch
import torch.nn as nn
import torch.nn.functional as F


class VGGbase(nn.Module):
    def __init__(self):
        super(VGGbase, self).__init__()

        # 卷积层，输入通道数为 3（表示 RGB 彩色图像），输出通道数为 64，卷积核大小为 3x3，填充为 1，步长为 1。
        # 填充是指在输入图像周围添加额外的像素值，以使得卷积核在图像边缘也能够进行有效的卷积操作。
        # 在这里，填充为 1 意味着在输入图像的周围各添加一行或一列的像素值。
        # 输入通道数：对于图像数据，通常情况下，输入通道数等于图像的通道数。例如，对于 RGB 彩色图像，输入通道数为 3；对于灰度图像，输入通道数为 1。
        # 输出通道数: 输出通道数的选择通常是根据任务的复杂度和模型的层次结构来确定的。
        # 通常，随着网络深度的增加，输出通道数也会增加，这有助于网络学习到更加抽象和复杂的特征。
        # 对于全连接层或者输出层，输出通道数通常与任务的类别数相关联。例如，对于分类任务，输出通道数通常等于类别数。
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, 3, 1, 1),
            # 批标准化层，有助于提高模型收敛
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        # 最大池化层:  2x2 大小的最大池化操作，步长为 2。
        self.max_pooling1 = nn.MaxPool2d(2, 2)

        #
        self.conv2_1 = nn.Sequential(
            nn.Conv2d(64, 128, 3, 1, 1),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )

        self.conv2_2 = nn.Sequential(
            nn.Conv2d(128, 128, 3, 1, 1),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )

        self.max_pooling2 = nn.MaxPool2d(2, 2)

        self.conv3_1 = nn.Sequential(
            nn.Conv2d(128, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )

        self.conv3_2 = nn.Sequential(
            nn.Conv2d(256, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )

        self.max_pooling3 = nn.MaxPool2d(2, 2, 1)

        self.conv4_1 = nn.Sequential(
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.BatchNorm2d(512),
            nn.ReLU()
        )

        self.conv4_2 = nn.Sequential(
            nn.Conv2d(512, 512, 3, 1, 1),
            nn.BatchNorm2d(512),
            nn.ReLU()
        )
        self.max_pooling4 = nn.MaxPool2d(2, 2)

        # batchsize * 512 *2 * 2 --> batchsize * (512 * 4)
        # 将卷积层输出的特征展平后输入到一个全连接层，将输入的特征向量维度从 512x2x2 转换为 10 维（因为这里是一个分类任务，输出类别数为 10）
        self.fc = nn.Linear(512 * 4, 10)

    def forward(self, x):
        batchsize = x.size(0)
        out = self.conv1(x)
        out = self.max_pooling1(out)
        out = self.conv2_1(out)
        out = self.conv2_2(out)
        out = self.max_pooling2(out)
        out = self.conv3_1(out)
        out = self.conv3_2(out)
        out = self.max_pooling3(out)
        out = self.conv4_1(out)
        out = self.conv4_2(out)
        out = self.max_pooling4(out)

        out = out.view(batchsize, -1)
        out = self.fc(out)
        out = F.log_softmax(out, dim=1)

        return out


def VGGNet():
    return VGGbase()
